import gradio as gr
import joblib
import numpy as np
from PIL import Image

# 加载模型
model_path = 'models/best_model.joblib'
model = joblib.load(model_path)

def preprocess_image(image_array):
    """预处理图片"""
    image = Image.fromarray(image_array).resize((64, 64)).convert('L')
    return np.array(image).flatten() / 255.0

def predict_image(image_array):
    """预测图片类别"""
    processed_image = preprocess_image(image_array)
    prediction = model.predict([processed_image])[0]
    try:
        probabilities = model.predict_proba([processed_image])[0]
        confidence = max(probabilities)
    except AttributeError:
        confidence = 1.0  # 如果模型不支持概率预测，则置信度为1

    result = "猫" if prediction == 0 else "狗"
    return f"{result} (置信度: {confidence:.2%})"

# 创建Gradio界面
iface = gr.Interface(
    fn=predict_image,
    inputs=gr.Image(label="上传图片"),
    outputs=gr.Text(label="预测结果"),
    title="猫狗图片分类器",
    description="上传一张图像来预测它是猫还是狗。",
    examples=[["example_cat.jpg"], ["example_dog.jpg"]]
)

# 启动应用
if __name__ == "__main__":
    iface.launch(share=True)